/*
 * 多进程结构主进程库
 */

package parent

import (
	"errors"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"sieve_engine/lib/app"
	"sieve_engine/lib/config"
	"sieve_engine/lib/log"
	"sieve_engine/lib/xconn"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"syscall"
	"time"
)

var (
	fSock  = config.String("mp.parent.sock", "/run/"+filepath.Base(os.Args[0])+"-"+strconv.Itoa(os.Getpid())+".mp.sock", "mp parent unix socket")
	parent *Parent
)

type Parent struct {
	sock          string
	lock          sync.RWMutex
	children      map[string]*Child
	msgHandler    func(child *Child, msg []byte) []byte
	statusHandler func(child *Child, status int)
	envs          []string
}

func NewParent(sock string) *Parent {
	p := &Parent{sock: sock, lock: sync.RWMutex{}, children: map[string]*Child{}, statusHandler: func(child *Child, status int) {}}
	app.AsyncSafe(func() {
		<-app.CTX.Done()
		p.Update(map[string]string{})
	})
	return p
}

func init() {
	parent = NewParent(*fSock)
}

// 通信

func (p *Parent) listen() error {
	os.Remove(p.sock)
	l, err := net.Listen("unix", p.sock)
	if err != nil {
		return err
	}
	app.Async(func() {
		for {
			if conn, err := l.Accept(); err == nil {
				p.handleConn(conn)
			} else if app.CTX.Err() != nil {
				return
			}
		}
	})
	app.AsyncSafe(func() {
		<-app.CTX.Done()
		l.Close()
		os.Remove(p.sock)
	})
	return nil
}

func (p *Parent) getChildByPid(pid int) *Child {
	p.lock.RLock()
	defer p.lock.RUnlock()
	for _, c := range p.children {
		if c.Pid() == pid {
			return c
		}
	}
	return nil
}

func (p *Parent) handleConn(conn net.Conn) {
	app.Async(func() {
		var ch *Child
		c, err := xconn.NewConn(conn, &xconn.Config{OnInit: func(rw xconn.RawRW) error {
			b, err := rw.Read()
			if err != nil {
				return err
			}
			pid, err := strconv.Atoi(string(b))
			if err != nil {
				return err
			}
			ch = p.getChildByPid(pid)
			if ch == nil {
				return errors.New("child not found")
			}
			return nil
		}})
		if err != nil {
			log.Err(err)
			return
		}
		c.Handle(func(data []byte) []byte {
			return p.msgHandler(ch, data)
		})
		ch.conn = c
		p.statusHandler(ch, ChildConnect)
	})
}

func (p *Parent) getChild(id string) *Child {
	p.lock.RLock()
	defer p.lock.RUnlock()
	return p.children[id]
}

func (p *Parent) HandleMessage(handler func(child *Child, req []byte) []byte) error {
	p.msgHandler = handler
	return p.listen()
}

func HandleMessage(handler func(child *Child, req []byte) []byte) error {
	return parent.HandleMessage(handler)
}

func (p *Parent) Request(id string, data []byte) ([]byte, error) {
	c := p.getChild(id)
	if c == nil {
		return nil, errors.New("child not found: " + id)
	}
	return c.Request(data)
}

func Request(id string, data []byte) ([]byte, error) {
	return parent.Request(id, data)
}

func (p *Parent) Post(id string, msg []byte) error {
	c := p.getChild(id)
	if c == nil {
		return errors.New("child not found: " + id)
	}
	return c.Post(msg)
}

func Post(id string, msg []byte) error {
	return parent.Post(id, msg)
}

func (p *Parent) AddEnv(env string) {
	p.envs = append(p.envs, env)
}

func AddEnv(env string) {
	parent.AddEnv(env)
}

// 进程

const (
	ChildStop    = 0
	ChildPrepare = 1
	ChildRun     = 2
	ChildConnect = 3
)

func (p *Parent) Update(list map[string]string) {
	p.lock.Lock()
	defer p.lock.Unlock()

	ids := []string{}
	for id := range p.children {
		ids = append(ids, id)
	}

	for _, id := range ids {
		if _, ok := list[id]; !ok {
			p.children[id].Stop()
			delete(p.children, id)
		}
	}
	for id, cmd := range list {
		if c, ok := p.children[id]; ok {
			if c.cmd != cmd {
				c.cmd = cmd
				c.Kill()
			}
		} else {
			c := &Child{parent: p, id: id, cmd: cmd, checkPerm: os.FileMode(0755)}
			p.children[id] = c
			c.Start()
		}
	}
}

func Update(list map[string]string) {
	parent.Update(list)
}

func (p *Parent) HandleStatus(handler func(child *Child, status int)) {
	p.statusHandler = handler
}

func HandleStatus(handler func(child *Child, status int)) {
	parent.HandleStatus(handler)
}

// /

type Child struct {
	parent    *Parent
	id        string
	cmd       string
	process   *os.Process
	conn      *xconn.Conn
	disabled  bool
	starting  int32
	checkPerm os.FileMode
}

func (c *Child) Id() string {
	return c.id
}

func (c *Child) Cmd() string {
	return c.cmd
}

func (c *Child) Start() {
	a := atomic.AddInt32(&c.starting, 1)
	defer atomic.AddInt32(&c.starting, -1)
	if a != 1 {
		return
	}

	c.disabled = false

	app.AsyncSafe(c.loop)
}

func (c *Child) loop() {
	for {
		s := strings.Fields(c.cmd)
		cmd := exec.Command(s[0])
		if len(s) > 1 {
			cmd = exec.Command(s[0], s[1:]...)
		}
		cmd.Env = append(os.Environ(), "__MP_CHILD_ID__="+c.id)
		if c.parent.msgHandler != nil && c.parent.sock != "" {
			cmd.Env = append(cmd.Env, "__MP_PARENT_SOCK__="+c.parent.sock)
		}
		cmd.Env = append(cmd.Env, "__DAEMON_IGNORED__=1")
		cmd.Env = append(cmd.Env, "__SYSTEMD_IGNORED__=1")
		cmd.Env = append(cmd.Env, "__CONFIG_DUMP__=1")
		for _, env := range c.parent.envs {
			cmd.Env = append(cmd.Env, env)
		}
		c.process = nil

		if c.disabled || app.CTX.Err() != nil {
			break
		}

		c.parent.statusHandler(c, ChildPrepare)
		c.checkAndSetExecPerm()

		if err := cmd.Start(); err == nil {
			c.process = cmd.Process
			c.parent.statusHandler(c, ChildRun)

			cmd.Wait()

			c.process = nil
			if c.conn != nil {
				c.conn.Close()
				c.conn = nil
			}
			c.parent.statusHandler(c, ChildStop)
		} else {
			log.Err(err)
		}

		if c.disabled || app.CTX.Err() != nil {
			break
		}

		time.Sleep(time.Second)
	}
}

func (c *Child) checkAndSetExecPerm() {
	if c.checkPerm == 0 {
		return
	}
	s := strings.Fields(c.cmd)
	f, err := os.Stat(s[0])
	if err != nil {
		log.Err(err)
		return
	}
	if c.checkPerm != f.Mode().Perm() {
		if err := os.Chmod(s[0], c.checkPerm); err != nil {
			log.Err(err)
		}
	}
}

func (c *Child) Pid() int {
	if proc := c.process; proc != nil {
		return proc.Pid
	}
	return 0
}

func (c *Child) IsRunning() bool {
	return c.Pid() > 0
}

func (c *Child) Stop() {
	c.disabled = true
	c.Kill()
}

func (c *Child) Kill() {
	if c.process != nil {
		c.process.Signal(syscall.SIGTERM)
		c.process = nil
	}
}

func (c *Child) Reload() {
	if c.process != nil {
		c.process.Signal(syscall.SIGHUP)
	}
}

func (c *Child) Post(msg []byte) error {
	if c.conn == nil {
		return errors.New("child has no connection: " + c.id)
	}
	return c.conn.Post(msg)
}

func (c *Child) Request(data []byte) ([]byte, error) {
	if c.conn == nil {
		return nil, errors.New("child has no connection: " + c.id)
	}
	return c.conn.Request(data)
}

func (c *Child) SetCheckPerm(perm os.FileMode) {
	c.checkPerm = perm
}
